/**
 * Created with IntelliJ IDEA.
 * Description:
 * User: syyyy
 * Date: 2025-10-25
 * Time: 10:47
 */
import java.util.Arrays;

public class StrassenMatrixMultiplication {

    // 最小拆分阈值（小于该值时用常规乘法，避免递归开销）
    private static final int THRESHOLD = 64;

    /**
     * Strassen算法入口：计算两个n阶矩阵的乘积
     * @param A 第一个n阶矩阵
     * @param B 第二个n阶矩阵
     * @return 乘积矩阵C = A * B
     */
    public static int[][] multiply(int[][] A, int[][] B) {
        int n = A.length;

        // 检查矩阵合法性（必须为方阵且维度匹配）
        if (n != A[0].length || n != B.length || n != B[0].length) {
            throw new IllegalArgumentException("矩阵必须为同阶方阵");
        }

        // 若矩阵规模小于阈值，直接用常规乘法
        if (n <= THRESHOLD) {
            return conventionalMultiply(A, B);
        }

        // 确保矩阵阶数为2的幂（补0至最近的2的幂，简化拆分）
        int newN = nextPowerOfTwo(n);
        int[][] A_padded = padMatrix(A, newN);
        int[][] B_padded = padMatrix(B, newN);

        // 递归计算乘积（填充后的矩阵）
        int[][] C_padded = strassenRecursive(A_padded, B_padded, newN);

        // 裁剪掉填充的0，得到n阶结果
        int[][] C = new int[n][n];
        for (int i = 0; i < n; i++) {
            System.arraycopy(C_padded[i], 0, C[i], 0, n);
        }
        return C;
    }

    /**
     * 递归实现Strassen算法
     * @param A 矩阵A（阶数为2的幂）
     * @param B 矩阵B（阶数为2的幂）
     * @param n 矩阵阶数
     * @return 乘积矩阵C = A * B
     */
    private static int[][] strassenRecursive(int[][] A, int[][] B, int n) {
        // 递归终止条件：小规模矩阵用常规乘法
        if (n <= THRESHOLD) {
            return conventionalMultiply(A, B);
        }

        int half = n / 2;

        // 拆分矩阵A为4个子矩阵
        int[][] A11 = new int[half][half];
        int[][] A12 = new int[half][half];
        int[][] A21 = new int[half][half];
        int[][] A22 = new int[half][half];
        splitMatrix(A, A11, A12, A21, A22, half);

        // 拆分矩阵B为4个子矩阵
        int[][] B11 = new int[half][half];
        int[][] B12 = new int[half][half];
        int[][] B21 = new int[half][half];
        int[][] B22 = new int[half][half];
        splitMatrix(B, B11, B12, B21, B22, half);

        // 计算7个中间矩阵（核心优化：仅7次乘法）
        int[][] M1 = strassenRecursive(add(A11, A22), add(B11, B22), half);
        int[][] M2 = strassenRecursive(add(A21, A22), B11, half);
        int[][] M3 = strassenRecursive(A11, subtract(B12, B22), half);
        int[][] M4 = strassenRecursive(A22, subtract(B21, B11), half);
        int[][] M5 = strassenRecursive(add(A11, A12), B22, half);
        int[][] M6 = strassenRecursive(subtract(A21, A11), add(B11, B12), half);
        int[][] M7 = strassenRecursive(subtract(A12, A22), add(B21, B22), half);

        // 计算结果矩阵的4个子矩阵
        int[][] C11 = add(subtract(add(M1, M4), M5), M7);
        int[][] C12 = add(M3, M5);
        int[][] C21 = add(M2, M4);
        int[][] C22 = add(subtract(add(M1, M3), M2), M6);

        // 合并子矩阵为完整结果矩阵
        int[][] C = new int[n][n];
        mergeMatrix(C11, C12, C21, C22, C, half);

        return C;
    }

    /**
     * 常规矩阵乘法（三重循环，用于小规模矩阵）
     */
    private static int[][] conventionalMultiply(int[][] A, int[][] B) {
        int n = A.length;
        int[][] C = new int[n][n];
        for (int i = 0; i < n; i++) {
            for (int k = 0; k < n; k++) {
                int a = A[i][k];
                for (int j = 0; j < n; j++) {
                    C[i][j] += a * B[k][j];
                }
            }
        }
        return C;
    }

    /**
     * 矩阵加法：C = A + B
     */
    private static int[][] add(int[][] A, int[][] B) {
        int n = A.length;
        int[][] C = new int[n][n];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                C[i][j] = A[i][j] + B[i][j];
            }
        }
        return C;
    }

    /**
     * 矩阵减法：C = A - B
     */
    private static int[][] subtract(int[][] A, int[][] B) {
        int n = A.length;
        int[][] C = new int[n][n];
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                C[i][j] = A[i][j] - B[i][j];
            }
        }
        return C;
    }

    /**
     * 拆分矩阵为4个子矩阵
     * @param parent 原矩阵
     * @param c11 左上子矩阵
     * @param c12 右上子矩阵
     * @param c21 左下子矩阵
     * @param c22 右下子矩阵
     * @param half 子矩阵阶数（原矩阵阶数的一半）
     */
    private static void splitMatrix(int[][] parent, int[][] c11, int[][] c12,
                                    int[][] c21, int[][] c22, int half) {
        for (int i = 0; i < half; i++) {
            System.arraycopy(parent[i], 0, c11[i], 0, half);
            System.arraycopy(parent[i], half, c12[i], 0, half);
            System.arraycopy(parent[i + half], 0, c21[i], 0, half);
            System.arraycopy(parent[i + half], half, c22[i], 0, half);
        }
    }

    /**
     * 合并4个子矩阵为原矩阵
     */
    private static void mergeMatrix(int[][] c11, int[][] c12, int[][] c21, int[][] c22,
                                    int[][] parent, int half) {
        for (int i = 0; i < half; i++) {
            System.arraycopy(c11[i], 0, parent[i], 0, half);
            System.arraycopy(c12[i], 0, parent[i], half, half);
            System.arraycopy(c21[i], 0, parent[i + half], 0, half);
            System.arraycopy(c22[i], 0, parent[i + half], half, half);
        }
    }

    /**
     * 计算大于等于n的最小2的幂（确保矩阵可被均匀拆分）
     */
    private static int nextPowerOfTwo(int n) {
        if (n <= 0) return 1;
        int power = 1;
        while (power < n) {
            power <<= 1; // 等价于 power *= 2
        }
        return power;
    }

    /**
     * 填充矩阵至指定阶数（不足部分补0）
     */
    private static int[][] padMatrix(int[][] matrix, int newN) {
        int n = matrix.length;
        int[][] padded = new int[newN][newN];
        for (int i = 0; i < n; i++) {
            System.arraycopy(matrix[i], 0, padded[i], 0, n);
        }
        return padded;
    }

    /**
     * 打印矩阵（辅助函数）
     */
    public static void printMatrix(int[][] matrix) {
        for (int[] row : matrix) {
            System.out.println(Arrays.toString(row));
        }
    }

    // 测试示例
    public static void main(String[] args) {
        // 2阶矩阵测试
        int[][] A = {{1, 2}, {3, 4}};
        int[][] B = {{5, 6}, {7, 8}};
        System.out.println("矩阵A：");
        printMatrix(A);
        System.out.println("矩阵B：");
        printMatrix(B);

        int[][] C = multiply(A, B);
        System.out.println("乘积矩阵C = A * B：");
        printMatrix(C);
        // 预期输出：
        // [19, 22]
        // [43, 50]
    }
}